import math
from utils import invmod,generatePrime

class RSA:
    def __init__(self) -> None:
        pass


    def encrypt(self, pt: int, pk: tuple) -> int:
        """
        使用RSA算法和公钥对加密明文
        返回 C
        @C是加密结果
        """
        M, e, n = pt, pk[0], pk[1]
        C = pow(M,e,n)
        return C         


    def decrypt(self, ct: int, sk: tuple) -> int:
        """
        使用RSA算法和私钥对解密密文
        返回 M
        @M是解密结果
        """
        C, d, n = ct, sk[0], sk[1]
        M = pow(C,d,n)
        return M

    
    def generateKeys(self, size: int) -> tuple:
        """
        随机生成RSA算法的公钥对和私钥对
        返回 (e, n) 和 (d, n)
        """
        p = generatePrime(size)
        q = generatePrime(size)

        while p == q:
            p, q = self.generatePrime(6), self.generatePrime(6)

        n = p * q
        s = (p - 1) * (q - 1)
        e = self.generateE(s)
        d = self.generateD(e,s)

        return (e,n), (d,n)


    def generateE(self, s: int) -> int:
        """
        生成RSA算法中的公钥e
        返回 e
        @e是小于s的最小正质数
        """
        for e in range(2,s):
            if math.gcd(e,s) == 1:
                return e
    

    def generateD(self, e: int, s :int) -> int:
        """
        生成RSA算法中的私钥d
        返回 d
        (@d * @e) mod @s = 1
        """
        return invmod(e,s)

        

 
